SpaceToBatchND

将输入张量在空间维（height、width）按 block_shape 分块，并重新排布到 batch 维；同时按 paddings 在空间维两侧做零填充。输入/输出均为 NHWC 格式： [batch, height, width, channel]。

与 SpaceToBatch 计算逻辑相同；schema 中分块参数名为 block_shape （SpaceToBatch 中为 block_size）。

设输入形状为 \([N, H, W, C]\)，分块为 \([B_h, B_w]\)，填充为 \([P_{top}, P_{bottom}, P_{left}, P_{right}]\)，则输出形状为：

\[\begin{split}\begin{aligned} N_{out} &= N \cdot B_h \cdot B_w \\ H_{out} &= (H + P_{top} + P_{bottom}) / B_h \\ W_{out} &= (W + P_{left} + P_{right}) / B_w \\ C_{out} &= C \end{aligned}\end{split}\]

要求 \((H + P_{top} + P_{bottom})\) 能被 \(B_h\) 整除， \((W + P_{left} + P_{right})\) 能被 \(B_w\) 整除。

映射关系

对输出坐标 \((n_{out}, h_{out}, w_{out}, c)\)：

\[\begin{split}\begin{aligned} n_{in} &= n_{out} \bmod N \\ s_w &= \lfloor n_{out} / N \rfloor \bmod B_w \\ s_h &= \lfloor \lfloor n_{out} / N \rfloor / B_w \rfloor \\ h_{in} &= h_{out} \cdot B_h + s_h - P_{top} \\ w_{in} &= w_{out} \cdot B_w + s_w - P_{left} \end{aligned}\end{split}\]

若 \(h_{in}\)、\(w_{in}\) 落在输入有效范围内，则 \(O[n_{out}, h_{out}, w_{out}, c] = I[n_{in}, h_{in}, w_{in}, c]\)；否则写 0（padding 区域）。

输入：

input - 输入数据地址
block - 分块因子（schema: block_shape），格式为 [block_h, block_w]
paddings - 填充参数，格式为 [top, bottom, left, right]
in_shape - 输入形状，格式为 [batch, height, width, channel]
core_mask - 核掩码（仅共享存储版本使用）

输出：

output - 输出数据地址

支持平台：

FT78NE MT7004

备注

FT78NE 支持 int8, int16, int32, fp32, fp64, cplx64, cplx128
MT7004 支持 int16, int32, fp16, fp32, cplx64

共享存储版本:

void i8_space_to_batch_nd_s(const int8_t *input, int8_t *output, const int *block, const int *paddings, const int *in_shape, int core_mask)

void i16_space_to_batch_nd_s(const int16_t *input, int16_t *output, const int *block, const int *paddings, const int *in_shape, int core_mask)

void i32_space_to_batch_nd_s(const int *input, int *output, const int *block, const int *paddings, const int *in_shape, int core_mask)

void hp_space_to_batch_nd_s(const float16 *input, float16 *output, const int *block, const int *paddings, const int *in_shape, int core_mask)

void fp_space_to_batch_nd_s(const float *input, float *output, const int *block, const int *paddings, const int *in_shape, int core_mask)

void dp_space_to_batch_nd_s(const double *input, double *output, const int *block, const int *paddings, const int *in_shape, int core_mask)

void c64_space_to_batch_nd_s(const float *input, float *output, const int *block, const int *paddings, const int *in_shape, int core_mask)

void c128_space_to_batch_nd_s(const double *input, double *output, const int *block, const int *paddings, const int *in_shape, int core_mask)

C调用示例：

// MT7004 示例（共享存储多核，DDR 地址）
void TestSpaceToBatchNDSMCFp32(int core_mask) {
    int core_id = get_core_id();
    int logic_core_id = GetLogicCoreId(core_mask, core_id);
    int core_num = GetCoreNum(core_mask);
    float *input = (float *)0x81000000;
    float *output = (float *)0x82000000;
    int block[2] = {2, 2};
    int paddings[4] = {0, 0, 0, 0};
    int in_shape[4] = {1, 8, 8, 4};
    sys_bar(0, core_num);
    fp_space_to_batch_nd_s(input, output, block, paddings, in_shape, core_mask);
}

void main() {
    int core_mask = 0b1111;
    TestSpaceToBatchNDSMCFp32(core_mask);
}

私有存储版本:

void i8_space_to_batch_nd_p(const int8_t *input, int8_t *output, const int *block, const int *paddings, const int *in_shape)

void i16_space_to_batch_nd_p(const int16_t *input, int16_t *output, const int *block, const int *paddings, const int *in_shape)

void i32_space_to_batch_nd_p(const int *input, int *output, const int *block, const int *paddings, const int *in_shape)

void hp_space_to_batch_nd_p(const float16 *input, float16 *output, const int *block, const int *paddings, const int *in_shape)

void fp_space_to_batch_nd_p(const float *input, float *output, const int *block, const int *paddings, const int *in_shape)

void dp_space_to_batch_nd_p(const double *input, double *output, const int *block, const int *paddings, const int *in_shape)

void c64_space_to_batch_nd_p(const float *input, float *output, const int *block, const int *paddings, const int *in_shape)

void c128_space_to_batch_nd_p(const double *input, double *output, const int *block, const int *paddings, const int *in_shape)

C调用示例：

// MT7004 示例（私有存储单核，AM 地址）
void TestSpaceToBatchNDAMFp32(void) {
    float *input = (float *)0x10000000;
    float *output = (float *)0x10020000;
    int block[2] = {2, 2};
    int paddings[4] = {0, 0, 0, 0};
    int in_shape[4] = {1, 8, 8, 4};
    fp_space_to_batch_nd_p(input, output, block, paddings, in_shape);
}

void main() {
    TestSpaceToBatchNDAMFp32();
}